import time
from Ashare    import *
# 加载指标计算库
from MyTT import *
import talib
import pandas_ta as ta

'''
https://blog.csdn.net/popboy29/article/details/125826838
我们提取1000跟日K线数据，并给指标设定相同的参数，使用以上三种包分别计算指标。
比较速度结论
talib最快，MyTT其次,pandas_ta最差慢了3个数量级无法接受

'''
# ===============表格美化输出===============
def df_table(df,index):
    import prettytable as pt
    #利用prettytable对输出结果进行美化,index为索引列名:df_table(df,'market')
    tb = pt.PrettyTable()
    df = df.reset_index(drop = True)
    tb.add_column(index,df.index)
    for col in df.columns.values:#df.columns.values的意思是获取列的名称
        tb.add_column(col, df[col])
    print(tb)

def cal_indicators_MyTT(df_data):
    # MyTT 
    C = df_data['close']
    df_data['ma5'] = MA(C,5)
    df_data['bollUpper'],df_data['bollMiddle'],df_data['bollBottom']  = BOLL(C,N=20, P=2)
    df_data['rsi'] = RSI(C,14)
    df_data['diff'],df_data['dea'],df_data['macd']=MACD(C,SHORT=12,LONG=26,M=9)
    return df_data.round(3)
        
def cal_indicators_Talib(df_data):
    # talib
    C = df_data['close']
    df_data['ma5'] = talib.MA(C,5)
    df_data['bollUpper'],df_data['bollMiddle'],df_data['bollBottom'] = talib.BBANDS(C,timeperiod=20,nbdevup=2,nbdevdn=2,matype=0)
    df_data['rsi'] = talib.RSI(C, timeperiod=14)
    df_data['diff'], df_data['dea'], df_data['macd'] = talib.MACD(C, fastperiod=12, slowperiod=26, signalperiod=9)
    # macd值需要乘以2，才可以和国内的股票软件一致
    df_data['macd']=df_data['macd'].map(lambda x: x*2)
    return df_data.round(3)
    
def cal_indicators_Pandas_ta(df_data):
    # pandas-ta
    CustomStrategy = ta.Strategy(
        name="我的指标库",
        description="SMA 50, BBANDS, RSI, MACD",
        ta=[
            {"kind": "sma", "length": 5},
            {"kind": "bbands", "length": 20},
            {"kind": "rsi", "length": 14},
            {"kind": "macd", "fast": 12, "slow": 26, "signal":9},
        ]
    )

    # # To run your "Custom Strategy"
    # print('df.ta.cores',df_data.ta.cores)
    # df_data.ta.cores = 0
    df_data.ta.strategy(CustomStrategy)
    # macd值需要乘以2，才可以和国内的股票软件一致
    df_data['MACDh_12_26_9']=df_data['MACDh_12_26_9'].map(lambda x: x*2)
    return df_data.round(3)
    
if __name__ == "__main__":
    # 1.过Ashare   获取数据
    print('开始提取K线数据')
    # from  Ashare    import *
    from Ashare   import get_price
    df_data=get_price('sz300750',frequency='1d',count=1000)  #frequency='1d' 表是获取日K,count=1000，表示获取1000根K线
    # print('Ashare   行情获取\n',df_data)
    df_table(df_data.tail(20),'df_data')

    # 2. 测试计时开始，测试哪个就把if 后面的0改为1即可，其它改成0。
    time1 = time.time()
    print('开始计算指标')
    print('------cal_indicators_MyTT')
    df = cal_indicators_MyTT(df_data)
    df_table(df.tail(20),'MyTT')
    time2 = time.time()
    print("cal_indicators_MyTT计算指标耗时:",time2-time1,'秒') 
    # cal_indicators_MyTT 计算指标耗时: 0.022001266479492188 秒    
    # cal_indicators_MyTT计算指标耗时: 0.007207155227661133 秒

    print('------cal_indicators_Talib')
    df = cal_indicators_Talib(df_data)
    df_table(df.tail(20),'Talib')
    time3 = time.time()
    print("cal_indicators_Talib计算指标耗时:",time3-time2,'秒') 
    # cal_indicators_Talib 计算指标耗时: 0.019998788833618164 秒
    # cal_indicators_Talib计算指标耗时: 0.005984067916870117 秒

    print('------cal_indicators_Pandas_ta')
    df = cal_indicators_Pandas_ta(df_data)
    df_table(df.tail(20),'Pandas_ta')
    time4 = time.time()
    print("cal_indicators_Pandas_ta计算指标耗时:",time4-time3,'秒') 
    # cal_indicators_Pandas_ta 计算指标耗时: 2.5784144401550293 秒
    # cal_indicators_Pandas_ta计算指标耗时: 4.085425615310669 秒

